{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "817d5bc4-2bd2-4608-b41b-9a642b778aea",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Set the environment variable to avoid memory leak warning\n",
    "os.environ[\"OMP_NUM_THREADS\"] = \"2\"\n",
    "\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "from sklearn.cluster import KMeans\n",
    "from sklearn.datasets import make_moons\n",
    "from sklearn.datasets import make_blobs\n",
    "\n",
    "import time\n",
    "\n",
    "import cv2\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "color_list = ['tab:blue', 'tab:orange', 'tab:green', 'tab:red', 'tab:purple', 'tab:brown', 'tab:pink', 'tab:gray', 'tab:olive', 'tab:cyan']\n",
    "\n",
    "label_size = 18 # Label size\n",
    "ticklabel_size = 14 # Tick label size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "515c9696-6634-48c4-a765-dae17d96a4c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate 1-D data\n",
    "np.random.seed(42)\n",
    "x = np.random.rand(80)\n",
    "\n",
    "# Plot\n",
    "fig, ax = plt.subplots(figsize=(6,6))\n",
    "ax.scatter(x, x, edgecolor='black', facecolor='pink', linewidth=1, s=6**2)\n",
    "ax.tick_params(axis='both', which='major', labelsize=ticklabel_size) # Set tick label size\n",
    "\n",
    "# plt.savefig('1D_data_base.png', dpi=300) # Make figure clearer\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "250204aa-c592-49cf-9ee0-6097f4c19848",
   "metadata": {},
   "outputs": [],
   "source": [
    "# k-Means clustering\n",
    "k = 2\n",
    "\n",
    "# Initialize cluster centers\n",
    "cluster_centers = [0.1, 0.2]\n",
    "\n",
    "# Initialize distance array\n",
    "distance = np.zeros((len(x), k))\n",
    "\n",
    "# Intercation\n",
    "max_iter = 8\n",
    "for iter_id in range(max_iter):\n",
    "    # Calculate distance between points to each center\n",
    "    for i in range(k):\n",
    "        distance[:,i] = np.abs(x - cluster_centers[i])\n",
    "\n",
    "    # Assign to closest centroid\n",
    "    cluster_idx = np.argmin(distance, axis=1)\n",
    "\n",
    "    # Display process\n",
    "    fig, ax = plt.subplots(figsize=(6,6))\n",
    "\n",
    "    # Display points close to center 1\n",
    "    ax.scatter(x[cluster_idx == 0], x[cluster_idx == 0], edgecolor='black', facecolor=color_list[8], linewidth=1, s=6**2)\n",
    "    ax.axvline(x=cluster_centers[0], color=color_list[8], linestyle='--', linewidth=1)\n",
    "\n",
    "    # Display points close to center 2\n",
    "    ax.scatter(x[cluster_idx == 1], x[cluster_idx == 1], edgecolor='black', facecolor=color_list[6], linewidth=1, s=6**2)\n",
    "    ax.axvline(x=cluster_centers[1], color=color_list[6], linestyle='--', linewidth=1)\n",
    "\n",
    "    ax.tick_params(axis='both', which='major', labelsize=ticklabel_size) # Set tick label size\n",
    "    ax.set_title(f'Interation: {iter_id}, C1: {cluster_centers[0]:.2f}, C2: {cluster_centers[1]:.2f}', fontsize=label_size)\n",
    "\n",
    "    # plt.savefig(f'1D_data_iter{iter_id}.png', dpi=300) # Make figure clearer\n",
    "    plt.show()\n",
    "\n",
    "    # Update cluster centers\n",
    "    for i in range(k):\n",
    "        cluster_centers[i] = np.mean(x[cluster_idx == i])\n",
    "\n",
    "# Display process\n",
    "fig, ax = plt.subplots(figsize=(6,6))\n",
    "\n",
    "# Display points close to center 1\n",
    "ax.scatter(x[cluster_idx == 0], x[cluster_idx == 0], edgecolor='black', facecolor=color_list[8], linewidth=1, s=6**2)\n",
    "ax.axvline(x=cluster_centers[0], color=color_list[8], linestyle='--', linewidth=1)\n",
    "\n",
    "# Display points close to center 2\n",
    "ax.scatter(x[cluster_idx == 1], x[cluster_idx == 1], edgecolor='black', facecolor=color_list[6], linewidth=1, s=6**2)\n",
    "ax.axvline(x=cluster_centers[1], color=color_list[6], linestyle='--', linewidth=1)\n",
    "\n",
    "ax.tick_params(axis='both', which='major', labelsize=ticklabel_size) # Set tick label size\n",
    "ax.set_title(f'Interation: {iter_id}, C1: {cluster_centers[0]:.2f}, C2: {cluster_centers[1]:.2f}', fontsize=label_size)\n",
    "\n",
    "# plt.savefig(f'1D_data_iter{iter_id}.png', dpi=300) # Make figure clearer\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3298d70-6093-4eb7-a902-efe43842cd1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the car number image\n",
    "image_carno = cv2.imread('D:/BaiduNetdiskDownload/car_num.jpg')\n",
    "\n",
    "# Convert from BGR to RGB\n",
    "image_carno = cv2.cvtColor(image_carno, cv2.COLOR_BGR2RGB)\n",
    "\n",
    "# Display the image\n",
    "fig, ax = plt.subplots(figsize=(7,7))\n",
    "img = ax.imshow(image_carno)\n",
    "plt.axis('off')  # Hide axes\n",
    "# plt.savefig('carno_base.png', dpi=300) # Make figure clearer\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "799b5fd1-f9bf-451f-aff9-7ec0b48d7aa1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Change image_carno into grey image\n",
    "image_carno_grey = cv2.cvtColor(image_carno, cv2.COLOR_BGR2GRAY)\n",
    "\n",
    "# Display the image\n",
    "fig, ax = plt.subplots(figsize=(7,7))\n",
    "img = ax.imshow(image_carno_grey, cmap='gray')\n",
    "plt.axis('off')  # Hide axes\n",
    "# plt.savefig('carno_grey.png', dpi=300) # Make figure clearer\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e21be06d-a91f-4c8e-b7a9-49631d86f20a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reshape image_carno_grey into a 1-D vector\n",
    "x = image_carno_grey.reshape(-1)\n",
    "\n",
    "# Using k-Means to separate background and foreground by pixels\n",
    "k = 2\n",
    "\n",
    "# Initialize cluster centers\n",
    "cluster_centers = np.random.rand(k) * 255\n",
    "start_centers = cluster_centers.copy()\n",
    "\n",
    "# Initialize distance array\n",
    "distance = np.zeros((len(x), k))\n",
    "\n",
    "# Intercation\n",
    "start_time = time.time()\n",
    "max_iter = 500\n",
    "for iter_id in range(max_iter):\n",
    "    # Calculate distance between points to each center\n",
    "    for i in range(k):\n",
    "        distance[:,i] = np.abs(x - cluster_centers[i])\n",
    "\n",
    "    # Assign to closest centroid\n",
    "    cluster_idx = np.argmin(distance, axis=1)\n",
    "\n",
    "    # Update cluster centers\n",
    "    cluster_centers_prior = cluster_centers.copy()\n",
    "    for i in range(k):\n",
    "        cluster_centers[i] = np.mean(x[cluster_idx == i])\n",
    "\n",
    "    # Check if cluster_centers are stable enough to stop training\n",
    "    print(f'Iteration {iter_id}: Updated centers {cluster_centers}, Prior centers {cluster_centers_prior}')\n",
    "    if np.sum(np.abs(cluster_centers-cluster_centers_prior)) == 0:\n",
    "        break\n",
    "\n",
    "    cluster_centers_prior = cluster_centers\n",
    "\n",
    "end_time = time.time()\n",
    "print(f'Stop after iteration {iter_id}, time consumption is {end_time-start_time}')\n",
    "\n",
    "# Generate segmentation mask\n",
    "carno_mask_pixel = np.zeros_like(cluster_idx)\n",
    "low_value_cluster = np.argmin(cluster_centers)\n",
    "carno_mask_pixel[cluster_idx != low_value_cluster] = 1 # Set pixels with higher grey value to 1\n",
    "carno_mask_pixel = carno_mask_pixel.reshape(image_carno_grey.shape)\n",
    "\n",
    "# Display the mask\n",
    "fig, ax = plt.subplots(figsize=(7,7))\n",
    "img = ax.imshow(carno_mask_pixel, cmap='gray')\n",
    "ax.set_title(f'Final Centers: {cluster_centers}, Iteration: {iter_id}', fontsize=label_size)\n",
    "plt.axis('off')  # Hide axes\n",
    "# plt.savefig('carno_mask_pixel_2.png', dpi=300) # Make figure clearer\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea62cabf-fc4a-4c6e-98ed-1171eb8db2bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the car number image\n",
    "image_bmwk = cv2.imread('D:/BaiduNetdiskDownload/bmwk.png')\n",
    "\n",
    "# Convert from BGR to RGB\n",
    "image_bmwk_rgb = cv2.cvtColor(image_bmwk, cv2.COLOR_BGR2RGB)\n",
    "\n",
    "x_r = image_bmwk_rgb[:, :, 0].reshape(-1) # Store colors in red channel\n",
    "x_g = image_bmwk_rgb[:, :, 1].reshape(-1) # Store colors in green channel\n",
    "x_b = image_bmwk_rgb[:, :, 2].reshape(-1) # Store colors in blue channel\n",
    "\n",
    "# Display the image with no margin\n",
    "plt.figure(figsize=(image_bmwk_rgb.shape[1]/100, image_bmwk_rgb.shape[0]/100))  # Convert pixels to inches\n",
    "plt.imshow(image_bmwk_rgb)\n",
    "plt.axis('off')  # Hide axes\n",
    "plt.subplots_adjust(left=0, right=1, top=1, bottom=0)  # Remove margins\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24f74211-3639-4d32-b765-6d0adb4fe19c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def KMeansImage(img, k):\n",
    "    # Get image size\n",
    "    w, h, c = img.shape\n",
    "\n",
    "    # Reshape image along channel\n",
    "    x = np.reshape(img, (w * h, c))\n",
    "\n",
    "    # Train k-Means model\n",
    "    mdl_km = KMeans(n_clusters=k, n_init='auto')\n",
    "    mdl_km.fit(x)\n",
    "\n",
    "    # Predict labels of each pixels\n",
    "    labels = mdl_km.predict(x).reshape(w, h)\n",
    "\n",
    "    # Get centers\n",
    "    center_colors = mdl_km.cluster_centers_ / 255.0\n",
    "\n",
    "    # Use center colors to generate compressed image\n",
    "    img_comp = np.zeros((w, h, c))\n",
    "    for i in range(w):\n",
    "        for j in range(h):\n",
    "            img_comp[i][j] = center_colors[labels[i][j]]\n",
    "    return img_comp, center_colors\n",
    "\n",
    "claster_num = [2, 4, 10, 20, 40, 80]\n",
    "for k in claster_num:\n",
    "    img_comp, center_colors = KMeansImage(image_bmwk_rgb, k)\n",
    "\n",
    "    # Display center colors\n",
    "    fig, ax = plt.subplots(figsize=(16,1))\n",
    "    ax.imshow([center_colors])\n",
    "    plt.axis('off')\n",
    "    # plt.savefig(f'bmwk_center_{k}.png', dpi=300)\n",
    "    plt.show()\n",
    "\n",
    "     # Display compressed image\n",
    "    plt.figure(figsize=(image_bmwk_rgb.shape[1]/100, image_bmwk_rgb.shape[0]/100))  # Convert pixels to inches\n",
    "    plt.imshow(img_comp)\n",
    "    plt.axis('off')\n",
    "    plt.subplots_adjust(left=0, right=1, top=1, bottom=0)  # Remove margins\n",
    "    # plt.savefig(f'bmwk_comp_{k}.png', format='png', bbox_inches='tight', pad_inches=0)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75a2204c-f856-447e-86fe-76a3aac04283",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Using make_blobs to generate data of ten clustering\n",
    "X_mb, y_mb = make_blobs(n_samples=800, n_features=2, centers=7, random_state=42)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(7,7))\n",
    "ax.scatter(X_mb[:, 0], X_mb[:, 1], marker=\"o\", c=y_mb, s=7**2, edgecolor=\"k\")\n",
    "plt.axis('off')\n",
    "# plt.savefig(f'make_blobs_base.png', dpi=300)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8992e72e-521c-4a9e-a9b2-5d6d2bae282f",
   "metadata": {},
   "outputs": [],
   "source": [
    "k_list = np.arange(2, 15, 1)\n",
    "sse_list = np.zeros(len(k_list))\n",
    "\n",
    "mdl_km_list = []\n",
    "for i in range(len(k_list)):\n",
    "    mdl_km = KMeans(n_clusters=k_list[i], n_init='auto')\n",
    "    mdl_km.fit(X_mb)\n",
    "    mdl_km_list.append(mdl_km)\n",
    "    sse_list[i] = mdl_km.inertia_\n",
    "\n",
    "# Plot sse_list\n",
    "fig, ax = plt.subplots(figsize=(8,6))\n",
    "ax.plot(k_list, sse_list, marker='o', linestyle='-', color='tab:blue')\n",
    "ax.set_xticks(k_list)\n",
    "ax.set_xlabel('Number of clusters (k)', fontsize=label_size)\n",
    "ax.set_ylabel('SSE', fontsize=label_size)\n",
    "ax.tick_params(axis='both', which='major', labelsize=ticklabel_size) # Set tick label size\n",
    "# plt.savefig(f'make_blobs_sse.png', dpi=300)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ab78c78-ea26-41f7-9ffe-e7349e618e76",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Display clustering result of k = 4,6,8\n",
    "k_disp = [4, 6, 8]\n",
    "for k in k_disp:\n",
    "    mdl_km = mdl_km_list[k-2]\n",
    "\n",
    "    fig, ax = plt.subplots(figsize=(8,8))\n",
    "    ax.scatter(X_mb[:, 0], X_mb[:, 1], marker=\"o\", c=mdl_km_list[k-2].labels_, s=8**2, edgecolor=\"k\")\n",
    "    ax.set_title(f'Number of clusters (k): {k}', fontsize=label_size)\n",
    "    plt.axis('off')\n",
    "\n",
    "    # plt.savefig(f'make_blobs_{k}.png', dpi=300)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abb454e2-be6a-474c-9415-a404b54183dd",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
